package com.gitee.checkpoint

import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.api.common.typeinfo.{TypeHint, TypeInformation}
import org.apache.flink.runtime.state.{FunctionInitializationContext, FunctionSnapshotContext}
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}

import scala.collection.mutable.ArrayBuffer

/**
  * 实现checkpoint功能
  * @param threshold
  */
class BufferingSink(val threshold: Int) extends RichSinkFunction[(String, Int)] with CheckpointedFunction {
  var checkpointedState: ListState[(String, Int)] = _
  var bufferedElements: ArrayBuffer[(String, Int)] = _

  def apply(threshold: Int) = {
    this.bufferedElements = new ArrayBuffer[(String, Int)]()
  }

  override def invoke(value: (String, Int), context: SinkFunction.Context[_]): Unit = {
    bufferedElements += value
    if (bufferedElements.size == threshold){
      for (element:(String,Int) <- bufferedElements){
        //send it to the sink
      }
      bufferedElements.clear()
    }
  }


  override def snapshotState(context: FunctionSnapshotContext): Unit = {
    /* 定期实现checkpoint
    *    就是把bufferedElements中的状态添加到checkpointedState
    *    */
    checkpointedState.clear()
    for (element:(String,Int) <- bufferedElements){
      checkpointedState.add(element)
    }
  }

  /* 恢复初始化的时候从保存的快照中获取数据，用于恢复到crash之前的状态
   *   就是把checkpointedState中的状态添加到bufferedElements  */
  override def initializeState(context: FunctionInitializationContext): Unit = {
    val descriptor = new ListStateDescriptor[(String, Int)]("buffered-elements", TypeInformation.of(new TypeHint[(String, Int)] {
    }))
    checkpointedState = context.getOperatorStateStore.getListState(descriptor)
    if(context.isRestored){
      /*for(element:(String,Int) <- checkpointedState.get()){
        bufferedElements += element
      }*/
    }
  }
}
